from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import mysql.connector
import os
from datetime import datetime
import uuid
import logging
import pwd
import grp
from ultralytics import YOLO  # 添加这行

app = Flask(__name__)
CORS(app, resources={r"/runs/detect/*": {"origins": "*"}})

# 设置日志
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

UPLOAD_FOLDER = '/www/wwwroot/bronzeai_backend/uploads'
PREDICT_FOLDER = '/www/wwwroot/bronzeai_backend/runs/detect'
YOLO_MODEL_PATH = '/www/wwwroot/bronzeai_backend/best.pt'

# MySQL配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'mf4B0zc4bFQGQQRcNUze',
    'database': 'yolov10_db'
}

# 获取web服务器用户和组
WEB_USER = 'www-data'  # 通常是 www-data，但可能需要根据您的设置进行调整
WEB_GROUP = 'www-data'

def set_permissions(path, is_dir=False):
    uid = pwd.getpwnam(WEB_USER).pw_uid
    gid = grp.getgrnam(WEB_GROUP).gr_gid
    
    os.chown(path, uid, gid)
    if is_dir:
        os.chmod(path, 0o755)
    else:
        os.chmod(path, 0o644)

# 创建目录并设置权限
for folder in [UPLOAD_FOLDER, PREDICT_FOLDER]:
    os.makedirs(folder, exist_ok=True)
    set_permissions(folder, is_dir=True)

# 设置静态文件夹的访问路径
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['PREDICT_FOLDER'] = PREDICT_FOLDER

# 全局变量来跟踪最新的预测文件夹
latest_predict_folder = None

def get_latest_predict_folder():
    predict_folders = [f for f in os.listdir(PREDICT_FOLDER) if f.startswith('predict')]
    if predict_folders:
        return max(predict_folders, key=lambda x: os.path.getmtime(os.path.join(PREDICT_FOLDER, x)))
    return None

@app.route('/runs/detect/<path:filename>')
def get_prediction_file(filename):
    logging.info(f"Attempting to serve file: {filename}")
    file_path = os.path.join(PREDICT_FOLDER, filename)
    if os.path.exists(file_path):
        return send_from_directory(PREDICT_FOLDER, filename)
    logging.error(f"File not found: {filename}")
    return jsonify({'error': 'File not found'}), 404

@app.route('/api2/upload', methods=['POST'])
def upload_image():
    logging.info("Received upload request")
    if 'image' not in request.files:
        logging.error("No image part in the request")
        return jsonify({'error': 'No image part'}), 400
    file = request.files['image']
    if file.filename == '':
        logging.error("No selected file")
        return jsonify({'error': 'No selected file'}), 400
    if file:
        filename = str(uuid.uuid4()) + '_' + datetime.now().strftime("%Y%m%d%H%M%S") + os.path.splitext(file.filename)[1]
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        set_permissions(filepath)
        logging.info(f"File saved and permissions set: {filepath}")
        
        try:
            # 保存文件路径到数据库
            conn = mysql.connector.connect(**db_config)
            cursor = conn.cursor()
            cursor.execute("INSERT INTO images (filename, filepath) VALUES (%s, %s)", (filename, filepath))
            conn.commit()
            cursor.close()
            conn.close()
            logging.info("File info saved to database")
        except mysql.connector.Error as err:
            logging.error(f"Database error: {err}")
            return jsonify({'error': 'Database error'}), 500
        
        return jsonify({'message': 'File uploaded successfully', 'filename': filename}), 200

@app.route('/api2/predict', methods=['POST'])
def predict():
    data = request.json
    filename = data.get('filename')
    if not filename:
        return jsonify({'error': 'No filename provided'}), 400

    filepath = os.path.join(UPLOAD_FOLDER, filename)
    if not os.path.exists(filepath):
        return jsonify({'error': 'File not found'}), 404

    try:
        # 加载模型
        model = YOLO(YOLO_MODEL_PATH)
    except Exception as e:
        return jsonify({'error': f'Error loading model: {str(e)}'}), 500

    try:
        # 运行预测并保存结果
        results = model.predict(source=filepath, save=True, save_txt=True, project=PREDICT_FOLDER)
    except Exception as e:
        return jsonify({'error': f'Error during prediction: {str(e)}'}), 500

    # 获取最新的预测结果文件夹
    predict_folders = [f for f in os.listdir(PREDICT_FOLDER) if f.startswith('predict')]
    latest_folder = max(predict_folders, key=lambda x: int(x[7:]) if x[7:].isdigit() else -1)
    result_path = os.path.join(PREDICT_FOLDER, latest_folder, filename)

    # 提取标签信息
    labels = []
    try:
        for result in results:
            for box in result.boxes:
                cls_id = int(box.cls.squeeze())  # 根据需要修改
                cls_name = model.names[cls_id]
                labels.append(cls_name)
        # 去除重复的标签
        labels = list(set(labels))
    except Exception as e:
        return jsonify({'error': f'Error processing results: {str(e)}'}), 500

    if os.path.exists(result_path):
        # 返回结果路径和标签信息
        return jsonify({'message': 'Prediction completed', 'result_path': f"/runs/detect/{latest_folder}/{filename}", 'labels': labels}), 200
    else:
        return jsonify({'error': 'Prediction result not found'}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7890, debug=True)